from datetime import datetime, timedelta
from typing import Dict, List, Optional

from openai import AsyncOpenAI
from sqlalchemy.orm import Session
from sqlalchemy import desc

from core.config import OPENAI_CONFIG
from models import engine, Chat, Message

MODEL = OPENAI_CONFIG[-1].get("model")


class Robot:
    def __init__(self, user_id: str):
        self.client = AsyncOpenAI(
            api_key=OPENAI_CONFIG[-1].get("api_key"),
            base_url=OPENAI_CONFIG[-1].get("base_url"),
        )
        self.session = Session(engine)
        self.messages: List = []
        self.situation: Optional[str] = None
        self.created_at: Optional[datetime] = datetime.now()
        self.chat_id: Optional[str] = None
        self._get_chat_and_messages(user_id=user_id)

    def _get_chat_and_messages(self, user_id: str):
        chat = (
                self.session.query(Chat)
                .filter(
                    Chat.user_id == user_id,
                    Chat.created_at >= datetime.now() - timedelta(hours=1),
                )
                .first()
        )
        if chat:
            self.chat_id = chat.id
        else:
            chat = Chat(user_id=user_id, created_at=datetime.now(), source="robot", mode=1)
            self.session.add(chat)
            self.session.commit()
            self.session.refresh(chat)
            self.chat_id = chat.id

        # Fetch messages for the chat
        latest_messages = (
            self.session.query(Message)
            .filter(Message.chat_id == self.chat_id)
            .order_by(desc(Message.index))
            .limit(10)
            .all()
        )
        for message in latest_messages:
            self.messages.append({
                "role": message.role,
                "content": message.content
            })
        self.messages.reverse()

        # TODO: Add System Prompt
        if chat.mode == 1:
            self.messages.insert(0, {
                "role": "system",
                "content": "You are a helpful assistant."
            })
        else:
            situation = chat.system_prompt
            self.messages.insert(0, {
                "role": "system",
                "content": "You are a professional robot."
            })


    def add_message(self, user_content: str):
        self.messages.append({
            "role": "user",
            "content": user_content
        })

    async def task_judgment(self):
        tools = [{
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get current temperature for a given location.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "City and country e.g. Bogotá, Colombia"
                        }
                    },
                    "required": [
                        "location"
                    ],
                    "additionalProperties": False
                },
                "strict": True
            }
        }]
        completion = await self.client.chat.completions.create(
            model=MODEL,
            messages=self.messages,
            tools=tool
        )
        if completion.choices[0].message.tool_calls:
            tool_call = completion.choices[0].message.tool_calls[0]
            return {
                "tool": tool_call.function.name,
                "args": tool_call.function.arguments
            }
        else:
            return None





